# -*-coding: utf-8 -*-
"""
    @Author : panjq
    @E-mail : pan_jinquan@163.com
    @Date   : 2021-03-12 10:05:31
"""
from detection.models.nn.net_retinaface_landm import RetinaFaceLandm
from detection.models.nn.net_slim import SlimLandm
from detection.models.nn.net_rfb import RFBLandm, RFB
from detection.models.ssd.model_slim_ssd import create_model_slim_ssd
from detection.models.ssd.model_rfb_ssd import create_model_rfb_ssd, create_model_rfb5s_ssd
from detection.models.ssd.model_rfb_landms_ssd import create_model_rfb_landm_ssd
from detection.models.ssd.model_mb_v2_ssd_ssd import create_mobilenetv2_ssd
from detection.models.ssd.model_mb_v2_landms_ssd import create_mobilenetv2_landm_ssd
from detection.models.ssd.model_yolox_ssd import create_model_yolox_v2_ssd, create_model_yolox_ssd
from detection.models.yolo.model_yolox import create_model_yolox


def build_ssd_model_v1(net_type: str, prior_boxes, width_mult=1.0, pretrained=True, phase='train', device="cuda:0"):
    if net_type.lower() == "mnet_landm".lower():
        # prior_cfg = mnet_face_config
        model = RetinaFaceLandm(prior_boxes=prior_boxes, phase=phase)
    elif net_type.lower() == "slim_landm".lower():
        model = SlimLandm(prior_boxes=prior_boxes, phase=phase)
    elif net_type.lower() == "RFB_landm".lower():
        # prior_cfg = rfb_face_config
        model = RFBLandm(prior_boxes=prior_boxes, phase=phase)
    elif net_type.lower() == "RFB".lower():
        # prior_cfg = rfb_face_config
        model = RFB(prior_boxes=prior_boxes, phase=phase)
    else:
        raise Exception("Error:{}".format(net_type))
    return model


def build_ssd_model_v2(net_type: str, prior_boxes, width_mult=1.0, pretrained=False, phase='train', device="cuda:0"):
    kwargs = {}
    if net_type.lower() == 'slim'.lower():
        create_model = create_model_slim_ssd
    elif net_type.lower() == 'RFB'.lower():
        create_model = create_model_rfb_ssd
    elif net_type.lower() == 'RFB5S'.lower():
        create_model = create_model_rfb5s_ssd
    elif net_type.lower() == 'mbv2'.lower():
        create_model = create_mobilenetv2_ssd
    elif net_type.lower() == 'RFB_Landm'.lower():
        create_model = create_model_rfb_landm_ssd
    elif net_type.lower() == 'MBV2_Landm'.lower():
        create_model = create_mobilenetv2_landm_ssd
    elif "yolox" in net_type.lower():
        create_model = create_model_yolox_ssd
        kwargs["net_type"] = net_type
    # elif "yolox" in net_type.lower():
    #     create_model = create_model_yolox_v2_ssd
    #     kwargs["net_type"] = net_type
    #     kwargs["depthwise"] = False
    #     kwargs["act"] = "silu"  # {"silu","relu","lrelu"}
    else:
        raise Exception("The net type is wrong.")
    is_test = not phase == 'train'
    model = create_model(prior_boxes=prior_boxes,
                         num_classes=prior_boxes.num_classes,
                         is_test=is_test,
                         width_mult=width_mult,
                         pretrained=pretrained,
                         device=device,
                         **kwargs,
                         )
    return model


def build_net_v3(net_type: str, num_classes, width_mult=1.0, pretrained=False, phase='train', device="cuda:0"):
    kwargs = {}
    if "yolox" in net_type.lower():
        create_model = create_model_yolox
        kwargs["depthwise"] = True
        kwargs["act"] = "silu"  # {"silu","relu","lrelu"}
    else:
        raise Exception("The net type is wrong.")
    is_test = not phase == 'train'
    model = create_model(net_type=net_type,
                         num_classes=num_classes,
                         width_mult=width_mult,
                         pretrained=pretrained,
                         is_test=is_test,
                         device=device,
                         **kwargs)
    return model


if __name__ == "__main__":
    import easydict
    import torch
    from basetrainer.utils import torch_tools
    from detection.anchor_utils.prior_box import PriorBox

    """
    ===RFB-1.0-default-[320,160]====
    Total memory: 22.97MB
    Total MAdd: 152.37MMAdd
    Total Flops: 78.36MFlops
    Total MemR+W: 49.27MB
    ===yolox-nano-0.5-yolox-[320,160]====
    Total memory: 9.69MB
    Total MAdd: 205.49MMAdd
    Total Flops: 102.74MFlops
    Total MemR+W: 16.87MB
    """
    # yolox-nano, yolox-tiny, yolox-s, yolox-m, yolox-l, yolox-x
    cfg = easydict.EasyDict({})
    cfg.input_size = [320, 160]
    cfg.net_type = "yolox-nano"
    # cfg.net_type = "RFB"
    cfg.priors_type = "default"
    cfg.width_mult = 1.0
    cfg.class_name = [0, 1]
    device = "cpu"
    pretrained = "/home/dm/data3/release/infrastructure/DMDetection/data/pretrained/pth/rfb_retails.pth"
    priorbox = PriorBox(cfg, input_size=cfg.input_size, priors_type=cfg.priors_type, freeze_header=True)
    # model = build_net_v2(cfg.net_type, priorbox, width_mult=cfg.width_mult, pretrained=pretrained, phase='train',
    #                      device=device)
    model = build_net_v3(cfg.net_type, num_classes=2, width_mult=cfg.width_mult, phase='test', device=device,pretrained=True)
    model = model.to(device)
    batch_size = 2
    inputs = torch.randn(size=(batch_size, 3, cfg.input_size[1], cfg.input_size[0]))
    inputs = inputs.to(device)
    model.eval()
    output = model(inputs)

    print("inputs.shape:{}".format(inputs.shape))
    # print("output.shape:{}".format(output.shape))
    file = "_".join([cfg.net_type, str(cfg.width_mult), str(cfg.input_size[0]), str(cfg.input_size[1])])
    # torch.save(model.state_dict(), file + ".pth")
    # torch_tools.summary_model(model, batch_size=1, input_size=cfg.input_size, device=device)
    # torch_tools.torchinfo_summary(model, batch_size=1, input_size=cfg.input_size, device=device)
    # summary(model, input_size=(1, input_size[1], input_size[0]), batch_size=batch_size, device=device)
    # stat(model, (1, input_size[0], input_size[1]))
    # print("===" * 10)
    # torch_tools.plot_model(model, output)
